from segment_anything import SamPredictor, sam_model_registry
import cv2
import numpy as np

model_type = "vit_h"  # 替换为实际使用的模型类型
checkpoint_path = '/home/JSDC/017254/code/gitee/map_learing/map_grid/04_xyzrpy_visual/demo/sam_vit_h_4b8939.pth'
sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
predictor = SamPredictor(sam)
img_path = '/home/JSDC/017254/code/gitee/map_learing/map_grid/04_xyzrpy_visual/demo/1744185603703480.jpg'
image = cv2.imread(img_path)
predictor.set_image(image)


# 输入提示
input_prompts_points = {
    "point_coords": np.array([(700, 300)]),  # 将点坐标转换为 NumPy 数组
    "point_labels": [1]  # 对应的标签，1表示目标对象，0表示背景
}

# 进行预测
masks_points, scores, logits = predictor.predict(
    point_coords=input_prompts_points["point_coords"],
    point_labels=input_prompts_points["point_labels"]
)

# 打印分割结果
print("Masks using points:", masks_points)
print("Scores:", scores)
print("Logits:", logits)


threshold = 128  # 选择合适的阈值，通常是 0 到 255 之间
# 可视化结果
for i, mask in enumerate(masks_points):


    scaled_mask = (mask * 255).astype('uint8')

        # 二值化处理
    _, binary_mask = cv2.threshold(scaled_mask, threshold, 255, cv2.THRESH_BINARY)

        # 将二值化掩码存储为黑白 PNG 文件
    cv2.imwrite(f'Binary_Mask_{i}.png', binary_mask)
    # 创建一个彩色掩码
    color_mask = (mask * 255).astype('uint8')
    color_mask = cv2.applyColorMap(color_mask, cv2.COLORMAP_JET)

    # 将掩码叠加到原始图像上
    overlay = cv2.addWeighted(image, 0.5, color_mask, 0.5, 0)

    # 显示结果
    # cv2.imshow(f'Mask {i}', overlay)
    cv2.imwrite(f'Mask {i}.jpg', overlay)
    # cv2.waitKey(0)

# cv2.destroyAllWindows()